import torch
import torch.nn.functional as F

Q = torch.tensor([[0.0000000, 1.0000000, 2.0000000, 3.0000000, 4.0000000, 5.0000000, 6.0000000, 7.0000000],
                  [0.0000000, 1.0000000, 2.0000000, 3.0000000,
                      4.0000000, 5.0000000, 6.0000000, 7.0000000],
                  [0.0000000, 1.0000000, 2.0000000, 3.0000000, 4.0000000, 5.0000000, 6.0000000, 7.0000000]])
Q = Q.view(1, 1, 3, 8)

K = Q.clone()
V = Q.clone()

out = F.scaled_dot_product_attention(Q, K, V)

print(out)
